package cn.jtool.dig.server.comm.real;

import cn.jtool.dig.common.model.*;
import cn.jtool.dig.common.netty.CtxUtils;
import cn.jtool.dig.common.netty.SNGenerator;
import cn.jtool.dig.common.netty.SessionIdGen;
import cn.jtool.dig.common.netty.SessionManager;
import cn.jtool.dig.server.comm.MessageExchange;
import cn.jtool.dig.server.comm.StageManager;
import cn.jtool.dig.server.dao.entity.LogConnect;
import cn.jtool.dig.server.dao.entity.PortMapping;
import cn.jtool.dig.server.service.LogService;
import cn.jtool.dig.server.service.PortMappingService;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import java.io.UnsupportedEncodingException;
import java.util.Arrays;
import java.util.Date;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

@Slf4j
@Component
@ChannelHandler.Sharable
public class RealConnHander extends ChannelInboundHandlerAdapter {

    @Resource
    private PortMappingService portMappingService;


    @Resource
    private LogService logService;

    @Override
    public void channelActive(ChannelHandlerContext ctx) {
        int port = CtxUtils.getLocalInetSocketAddress(ctx).getPort();
        PortMapping portMapping = this.portMappingService.selectBySourcePort(port);
        if(Objects.isNull(portMapping)){
            log.warn("[real]连接激活, 未找到转发记录, 连接端口: {}, 远程地址: {}", port, ctx.channel().remoteAddress());
        }else{
            boolean iptablesResult = CtxUtils.iptables(ctx,
                    StringUtils.isNotBlank(portMapping.getWhiteIps()) ? Arrays.stream(portMapping.getWhiteIps().split(",")).toList() : null,
                    StringUtils.isNotBlank(portMapping.getBlackIps()) ? Arrays.stream(portMapping.getBlackIps().split(",")).toList() : null);
            if(!iptablesResult){
                log.warn("[real] {} 非法ip接入, ip: {}", portMapping.getPortMappingName(), CtxUtils.getRemoteInetSocketAddress(ctx).getHostName());
                ctx.channel().close();
                return;
            }
            long sessionId = SessionIdGen.gen();
            CtxUtils.setAttribute(ctx, RealAttrKey.SESSION_ID, sessionId);
            CtxUtils.setAttribute(ctx, RealAttrKey.OUTWARD_DIG_CONF, portMapping);
            SessionManager.addSession(
                    Session.builder()
                            .id(sessionId)
                            .licenseId(portMapping.getLicenseId())
                            .portMappingId(portMapping.getId())
                            .sessionType(Session.SessionType.REAL)
                            .channelHandlerContext(ctx)
                            .receiveBytes(new AtomicLong(0))
                            .sendBytes(new AtomicLong(0))
                            .createTime(new Date())
                            .updateTime(new Date())
                            .build()
            );

            PacketHeader header = new PacketHeader();
            header.setCommand(Commands.CREATE_SESSION.getCode());
            header.setDigId(portMapping.getLicenseId());
            header.setChannelId(sessionId);
            header.setSn(SNGenerator.gen());

            PacketBody body = new PacketBody();
            try {
                byte[] ipBytes = portMapping.getTargetIp().getBytes("GBK");
                body.writeByte(ipBytes.length);
                body.write(ipBytes);
                body.writeWord(portMapping.getTargetPort());

                MessageExchange.upAdmin(Packet.builder().header(header).body(body).build());
            } catch (UnsupportedEncodingException e) {
                throw new RuntimeException(e);
            }
            ctx.executor().schedule(() -> {
                LogConnect logConnect = this.logService.saveLogConnect(
                        LogConnect.builder()
                                .licenseId(portMapping.getLicenseId())
                                .portMappingId(portMapping.getId())
                                .sourceServer(CtxUtils.getAddress(ctx))
                                .connectTime(new Date())
                                .recBytes(0L)
                                .sendBytes(0L)
                                .build()
                );
                CtxUtils.setAttribute(ctx, RealAttrKey.LOG_CONNECT, logConnect);
            }, 0, TimeUnit.MILLISECONDS);
        }
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        log.info("[real]连接断开, 链路: {}, 远程地址: {}", CtxUtils.getAttribute(ctx, RealAttrKey.SESSION_ID), CtxUtils.getAddress(ctx));
        Long sessionId = CtxUtils.getAttribute(ctx, RealAttrKey.SESSION_ID);
        if(Objects.isNull(sessionId)){
            return;
        }
        Session session = SessionManager.getSession(sessionId);
        PortMapping portMapping = CtxUtils.getAttribute(ctx, RealAttrKey.OUTWARD_DIG_CONF);
        SessionManager.removeSession(sessionId);
        StageManager.removeStage(sessionId);
        if(Objects.nonNull(portMapping)) {
            PacketHeader header = new PacketHeader();
            header.setCommand(Commands.CLOSE_SESSION.getCode());
            header.setDigId(portMapping.getLicenseId());
            header.setChannelId(sessionId);
            header.setSn(SNGenerator.gen());
            MessageExchange.upAdmin(Packet.builder().header(header).build());


            ctx.executor().schedule(() -> {
                LogConnect logConnect = CtxUtils.getAttribute(ctx, RealAttrKey.LOG_CONNECT);
                if(Objects.nonNull(logConnect)){
                    logConnect.setDisConnectTime(new Date());
                    logConnect.setRecBytes(session.getReceiveBytes());
                    logConnect.setSendBytes(session.getSendBytes());
                    logConnect.setRecBytesSegment(session.getTrafficBySegment(Session.TrafficDir.REC, true));
                    logConnect.setSendBytesSegment(session.getTrafficBySegment(Session.TrafficDir.SEND, true));
                    this.logService.saveLogConnect(logConnect);
                }
            }, 500, TimeUnit.MILLISECONDS);
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        log.error("[real]连接异常, 远程地址: {}", ctx.channel().remoteAddress(), cause);
        this.channelInactive(ctx);
    }
}
